"""
简化版多因子螺纹钢策略
不依赖talib，使用基础技术指标
"""

import numpy as np
import pandas as pd
from datetime import datetime
from tqsdk import TqApi, TargetPosTask
import logging


class SimpleMultiFactorStrategy:
    """
    简化版多因子策略
    
    技术指标:
    1. 双均线系统
    2. RSI
    3. 布林带
    4. 价格动量
    """
    
    def __init__(self, api, symbol="rb", **kwargs):
        self.api = api
        self.symbol = symbol
        self.i_symbol = f"KQ.i@SHFE.{symbol}"
        self.m_symbol = f"KQ.m@SHFE.{symbol}"
        
        # 技术指标参数
        self.fast_ma = kwargs.get('fast_ma', 10)
        self.slow_ma = kwargs.get('slow_ma', 30)
        self.rsi_period = kwargs.get('rsi_period', 14)
        self.rsi_oversold = kwargs.get('rsi_oversold', 30)
        self.rsi_overbought = kwargs.get('rsi_overbought', 70)
        self.bb_period = kwargs.get('bb_period', 20)
        self.bb_std = kwargs.get('bb_std', 2.0)
        
        # 信号参数
        self.min_signals = kwargs.get('min_signals', 2)
        self.trend_filter = kwargs.get('trend_filter', True)
        
        # 风控参数
        self.max_position = kwargs.get('max_position', 1)
        self.atr_multiplier = kwargs.get('atr_multiplier', 2.0)
        self.min_interval = kwargs.get('min_interval', 30)
        
        # 数据订阅
        self.klines = api.get_kline_serial(self.i_symbol, 15 * 60, data_length=200)
        
        # 交易对象
        self.main_symbol = api.get_quote(self.m_symbol).underlying_symbol
        self.target_pos = TargetPosTask(api, self.main_symbol)
        
        # 交易状态
        self.entry_price = 0.0
        self.stop_loss_price = 0.0
        self.last_trade_time = 0
        
        # 日志
        self.logger = logging.getLogger(f"simple_multi_factor.{symbol}")
        
        print(f"✅ 简化多因子策略初始化完成")
        print(f"   快慢均线: {self.fast_ma}/{self.slow_ma}")
        print(f"   RSI参数: 周期{self.rsi_period}, 区间[{self.rsi_oversold}, {self.rsi_overbought}]")
        print(f"   布林带: 周期{self.bb_period}, 标准差{self.bb_std}")
        print(f"   最少信号数: {self.min_signals}")
    
    def calculate_sma(self, data, period):
        """计算简单移动平均"""
        if len(data) < period:
            return np.full(len(data), np.nan)
        
        result = np.full(len(data), np.nan)
        for i in range(period - 1, len(data)):
            result[i] = np.mean(data[i - period + 1:i + 1])
        return result
    
    def calculate_rsi(self, data, period):
        """计算RSI"""
        if len(data) < period + 1:
            return np.full(len(data), 50.0)
        
        deltas = np.diff(data)
        gains = np.where(deltas > 0, deltas, 0)
        losses = np.where(deltas < 0, -deltas, 0)
        
        avg_gains = np.full(len(data), np.nan)
        avg_losses = np.full(len(data), np.nan)
        
        # 初始平均值
        if len(gains) >= period:
            avg_gains[period] = np.mean(gains[:period])
            avg_losses[period] = np.mean(losses[:period])
            
            # 指数移动平均
            for i in range(period + 1, len(data)):
                avg_gains[i] = (avg_gains[i-1] * (period - 1) + gains[i-1]) / period
                avg_losses[i] = (avg_losses[i-1] * (period - 1) + losses[i-1]) / period
        
        # 计算RSI
        rsi = np.full(len(data), 50.0)
        for i in range(period, len(data)):
            if avg_losses[i] != 0:
                rs = avg_gains[i] / avg_losses[i]
                rsi[i] = 100 - (100 / (1 + rs))
            else:
                rsi[i] = 100
        
        return rsi
    
    def calculate_bollinger_bands(self, data, period, std_dev):
        """计算布林带"""
        if len(data) < period:
            return np.full(len(data), np.nan), np.full(len(data), np.nan), np.full(len(data), np.nan)
        
        sma = self.calculate_sma(data, period)
        
        upper = np.full(len(data), np.nan)
        lower = np.full(len(data), np.nan)
        
        for i in range(period - 1, len(data)):
            std = np.std(data[i - period + 1:i + 1])
            upper[i] = sma[i] + std_dev * std
            lower[i] = sma[i] - std_dev * std
        
        return upper, sma, lower
    
    def calculate_atr(self, high, low, close, period=14):
        """计算ATR"""
        if len(high) < period + 1:
            return close[-1] * 0.02
        
        tr = []
        for i in range(1, len(high)):
            tr1 = high[i] - low[i]
            tr2 = abs(high[i] - close[i-1])
            tr3 = abs(low[i] - close[i-1])
            tr.append(max(tr1, tr2, tr3))
        
        if len(tr) >= period:
            return np.mean(tr[-period:])
        else:
            return close[-1] * 0.02
    
    def calculate_indicators(self):
        """计算所有技术指标"""
        try:
            if len(self.klines) < max(self.slow_ma, self.bb_period, self.rsi_period) + 10:
                return {}
            
            close = self.klines.close.values
            high = self.klines.high.values
            low = self.klines.low.values
            
            # 移动平均
            fast_ma = self.calculate_sma(close, self.fast_ma)
            slow_ma = self.calculate_sma(close, self.slow_ma)
            
            # RSI
            rsi = self.calculate_rsi(close, self.rsi_period)
            
            # 布林带
            bb_upper, bb_middle, bb_lower = self.calculate_bollinger_bands(close, self.bb_period, self.bb_std)
            
            # ATR
            atr = self.calculate_atr(high, low, close)
            
            return {
                'close': close[-1],
                'fast_ma': fast_ma[-1] if not np.isnan(fast_ma[-1]) else 0,
                'slow_ma': slow_ma[-1] if not np.isnan(slow_ma[-1]) else 0,
                'prev_fast_ma': fast_ma[-2] if len(fast_ma) > 1 and not np.isnan(fast_ma[-2]) else fast_ma[-1],
                'prev_slow_ma': slow_ma[-2] if len(slow_ma) > 1 and not np.isnan(slow_ma[-2]) else slow_ma[-1],
                'rsi': rsi[-1] if not np.isnan(rsi[-1]) else 50,
                'bb_upper': bb_upper[-1] if not np.isnan(bb_upper[-1]) else 0,
                'bb_middle': bb_middle[-1] if not np.isnan(bb_middle[-1]) else 0,
                'bb_lower': bb_lower[-1] if not np.isnan(bb_lower[-1]) else 0,
                'atr': atr
            }
            
        except Exception as e:
            self.logger.error(f"指标计算失败: {e}")
            return {}
    
    def generate_signals(self, indicators):
        """生成交易信号"""
        if not indicators:
            return {'long': 0, 'short': 0, 'signals': []}
        
        signals = []
        long_score = 0
        short_score = 0
        
        close = indicators['close']
        
        # 1. 双均线信号
        if indicators['fast_ma'] > 0 and indicators['slow_ma'] > 0:
            if indicators['fast_ma'] > indicators['slow_ma'] and close > indicators['fast_ma']:
                long_score += 1
                signals.append("均线多头")
            elif indicators['fast_ma'] < indicators['slow_ma'] and close < indicators['fast_ma']:
                short_score += 1
                signals.append("均线空头")
        
        # 2. RSI信号
        rsi = indicators['rsi']
        if rsi < self.rsi_oversold:
            long_score += 1
            signals.append(f"RSI超卖({rsi:.1f})")
        elif rsi > self.rsi_overbought:
            short_score += 1
            signals.append(f"RSI超买({rsi:.1f})")
        
        # 3. 布林带信号
        if indicators['bb_lower'] > 0 and close < indicators['bb_lower']:
            long_score += 1
            signals.append("布林带下轨支撑")
        elif indicators['bb_upper'] > 0 and close > indicators['bb_upper']:
            short_score += 1
            signals.append("布林带上轨阻力")
        
        # 4. 价格动量信号
        if len(self.klines) >= 5:
            recent_closes = self.klines.close.iloc[-5:].values
            if recent_closes[-1] > recent_closes[0]:  # 5期上涨
                long_score += 0.5
                signals.append("价格动量向上")
            elif recent_closes[-1] < recent_closes[0]:  # 5期下跌
                short_score += 0.5
                signals.append("价格动量向下")
        
        # 5. 趋势过滤
        if self.trend_filter and indicators['slow_ma'] > 0:
            if indicators['slow_ma'] > indicators['prev_slow_ma']:
                short_score = max(0, short_score - 1)
            elif indicators['slow_ma'] < indicators['prev_slow_ma']:
                long_score = max(0, long_score - 1)
        
        return {
            'long': long_score,
            'short': short_score,
            'signals': signals
        }
    
    def on_bar(self):
        """主循环"""
        bar_count = 0
        while True:
            self.api.wait_update()
            if self.api.is_changing(self.klines):
                bar_count += 1
                current_time = self.klines.datetime.iloc[-1]
                current_price = self.klines.close.iloc[-1]
                volume = self.klines.volume.iloc[-1]
                
                if bar_count % 50 == 0:
                    print(f"处理第{bar_count}根K线: 时间={current_time}, 价格={current_price:.1f}")
                
                if volume != 0:
                    self.run()
    
    def run(self):
        """策略执行逻辑"""
        try:
            # 更新主力合约
            self.main_symbol = self.api.get_quote(self.m_symbol).underlying_symbol
            self.target_pos = TargetPosTask(self.api, self.main_symbol)
            
            # 平掉非主力持仓
            positions = self.api.get_position()
            for sym, pos in positions.items():
                if sym != self.main_symbol and pos.pos != 0:
                    TargetPosTask(self.api, sym).set_target_volume(0)
            
            main_pos = positions.get(self.main_symbol, None)
            current_pos = main_pos.pos if main_pos else 0
            current_price = self.klines.close.iloc[-1]
            current_time = self.klines.datetime.iloc[-1]
            
            # 控制交易频率
            if current_time - self.last_trade_time < self.min_interval * 60 * 1000:
                return
            
            # 计算技术指标
            indicators = self.calculate_indicators()
            if not indicators:
                return
            
            # 生成交易信号
            signal_result = self.generate_signals(indicators)
            
            # 调试信息
            if hasattr(self, 'run_count'):
                self.run_count += 1
            else:
                self.run_count = 1
                
            if self.run_count % 40 == 0:
                print(f"简化多因子第{self.run_count}次: 价格={current_price:.1f}, "
                      f"多头={signal_result['long']:.1f}, 空头={signal_result['short']:.1f}, "
                      f"RSI={indicators['rsi']:.1f}, 持仓={current_pos}")
            
            if current_pos == 0:  # 无持仓
                if signal_result['long'] >= self.min_signals:
                    # 做多
                    self.entry_price = current_price
                    stop_distance = self.atr_multiplier * indicators['atr']
                    self.stop_loss_price = current_price - stop_distance
                    
                    self.target_pos.set_target_volume(self.max_position)
                    self.last_trade_time = current_time
                    
                    print(f"🟢 简化多因子做多: 价格={current_price:.1f}, 信号={signal_result['long']:.1f}, "
                          f"原因={signal_result['signals']}, 止损={self.stop_loss_price:.1f}")
                
                elif signal_result['short'] >= self.min_signals:
                    # 做空
                    self.entry_price = current_price
                    stop_distance = self.atr_multiplier * indicators['atr']
                    self.stop_loss_price = current_price + stop_distance
                    
                    self.target_pos.set_target_volume(-self.max_position)
                    self.last_trade_time = current_time
                    
                    print(f"🔴 简化多因子做空: 价格={current_price:.1f}, 信号={signal_result['short']:.1f}, "
                          f"原因={signal_result['signals']}, 止损={self.stop_loss_price:.1f}")
            
            else:  # 有持仓
                should_close = False
                close_reason = ""
                
                if current_pos > 0:  # 多头持仓
                    if current_price <= self.stop_loss_price:
                        should_close = True
                        close_reason = "多头止损"
                    elif signal_result['short'] >= self.min_signals:
                        should_close = True
                        close_reason = f"信号反转(空头{signal_result['short']:.1f})"
                
                elif current_pos < 0:  # 空头持仓
                    if current_price >= self.stop_loss_price:
                        should_close = True
                        close_reason = "空头止损"
                    elif signal_result['long'] >= self.min_signals:
                        should_close = True
                        close_reason = f"信号反转(多头{signal_result['long']:.1f})"
                
                if should_close:
                    self.target_pos.set_target_volume(0)
                    self.last_trade_time = current_time
                    print(f"⚪ {close_reason}: 价格={current_price:.1f}")
                    
        except Exception as e:
            self.logger.error(f"策略执行失败: {e}")
    
    def get_performance_metrics(self):
        """获取性能指标"""
        return {
            'note': '简化多因子策略',
            'strategy_type': '多技术指标组合策略(简化版)',
            'timeframe': '15分钟',
            'factors': ['双均线', 'RSI', '布林带', '价格动量']
        }